#!/usr/bin/env python
# -*- coding: utf-8 -*-

import yaml
import subprocess
import os
import sys
import time
import argparse
from pathlib import Path

class BatchServiceLauncher:
    def __init__(self, config_file=None, delay=2):
        """
        批量服务启动器
        
        Args:
            config_file (str): 配置文件路径，默认使用项目中的config.yaml
            delay (int): 每个服务启动之间的延迟时间（秒）
        """
        # 设置默认配置文件路径
        if config_file is None:
            # 获取当前脚本所在目录的父目录，然后找到config/config.yaml
            script_dir = Path(__file__).parent
            project_root = script_dir.parent
            self.config_file = project_root / "config" / "config.yaml"
        else:
            self.config_file = Path(config_file)
        
        self.delay = delay
        self.launched_screens = []
        
        # 验证配置文件是否存在
        if not self.config_file.exists():
            print(f"错误：配置文件 {self.config_file} 不存在")
            sys.exit(1)
        
        # 验证single_service_runner.py是否存在
        self.single_runner_path = script_dir / "single_service_runner.py"
        if not self.single_runner_path.exists():
            print(f"错误：single_service_runner.py 文件 {self.single_runner_path} 不存在")
            sys.exit(1)
        
        print(f"配置文件路径: {self.config_file}")
        print(f"单服务运行器路径: {self.single_runner_path}")
        print(f"服务启动间隔: {self.delay} 秒")
    
    def load_services(self):
        """从配置文件加载所有服务"""
        try:
            with open(self.config_file, 'r', encoding='utf-8') as f:
                config = yaml.safe_load(f)
            
            if not config or 'services' not in config:
                print(f"错误：配置文件 {self.config_file} 格式无效，缺少 'services' 键")
                sys.exit(1)
            
            services = config['services']
            if not services:
                print("警告：配置文件中没有找到任何服务")
                return []
            
            print(f"从配置文件中加载了 {len(services)} 个服务")
            return services
        
        except Exception as e:
            print(f"错误：加载配置文件时发生异常: {e}")
            sys.exit(1)
    
    def check_screen_installed(self):
        """检查系统是否安装了screen"""
        try:
            subprocess.run(['screen', '--version'], 
                         capture_output=True, check=True)
            return True
        except (subprocess.CalledProcessError, FileNotFoundError):
            print("错误：系统未安装screen，请先安装screen")
            print("Ubuntu/Debian: sudo apt-get install screen")
            print("CentOS/RHEL: sudo yum install screen")
            return False
    
    def get_screen_session_name(self, service_name):
        """为服务生成screen会话名称"""
        # 清理服务名称，确保它适合作为screen会话名
        clean_name = service_name.replace('/', '_').replace(' ', '_').replace('-', '_')
        return f"service_{clean_name}"
    
    def is_screen_session_running(self, session_name):
        """检查指定的screen会话是否正在运行"""
        try:
            result = subprocess.run(['screen', '-list'], 
                                  capture_output=True, text=True)
            return session_name in result.stdout
        except Exception:
            return False
    
    def launch_service_in_screen(self, service_config):
        """在screen会话中启动单个服务"""
        service_name = service_config.get('name')
        if not service_name:
            print("警告：跳过没有名称的服务配置")
            return False
        
        # 跳过空的script_path的服务（通常是停止服务的占位符）
        if service_config.get('type') == 'python' and not service_config.get('script_path'):
            print(f"跳过服务 '{service_name}' (script_path为空)")
            return False
        
        session_name = self.get_screen_session_name(service_name)
        
        # 检查会话是否已经存在
        if self.is_screen_session_running(session_name):
            print(f"警告：Screen会话 '{session_name}' 已存在，跳过服务 '{service_name}'")
            return False
        
        # 构建启动命令
        # 使用rosrun来启动single_service_runner.py
        launch_cmd = [
            'rosrun', 'service_launcher', 'single_service_runner.py', 
            service_name, 
            f'_config_file:={self.config_file}'
        ]
        
        # 构建screen命令
        screen_cmd = [
            'screen', '-dmS', session_name,  # -d -m 创建新会话并分离，-S 指定会话名
            'bash', '-c', 
            f"{' '.join(launch_cmd)}; echo '服务 {service_name} 已退出，按任意键关闭...'; read"
        ]
        
        try:
            print(f"正在启动服务 '{service_name}' 在screen会话 '{session_name}' 中...")
            print(f"执行命令: {' '.join(screen_cmd)}")
            
            # 启动screen会话
            subprocess.run(screen_cmd, check=True)
            
            # 等待一小段时间确保会话启动
            time.sleep(0.5)
            
            # 验证会话是否成功创建
            if self.is_screen_session_running(session_name):
                print(f"✓ 服务 '{service_name}' 已成功在screen会话 '{session_name}' 中启动")
                self.launched_screens.append(session_name)
                return True
            else:
                print(f"✗ 服务 '{service_name}' 启动失败，screen会话未找到")
                return False
                
        except subprocess.CalledProcessError as e:
            print(f"✗ 启动服务 '{service_name}' 时发生错误: {e}")
            return False
        except Exception as e:
            print(f"✗ 启动服务 '{service_name}' 时发生未知错误: {e}")
            return False
    
    def launch_all_services(self):
        """启动所有服务"""
        if not self.check_screen_installed():
            return False
        
        services = self.load_services()
        if not services:
            return True
        
        print(f"\n开始批量启动 {len(services)} 个服务...")
        print("=" * 50)
        
        success_count = 0
        failed_count = 0
        
        for i, service_config in enumerate(services, 1):
            service_name = service_config.get('name', f'未命名服务_{i}')
            print(f"\n[{i}/{len(services)}] 处理服务: {service_name}")
            
            if self.launch_service_in_screen(service_config):
                success_count += 1
            else:
                failed_count += 1
            
            # 在启动下一个服务前等待
            if i < len(services):  # 不是最后一个服务
                print(f"等待 {self.delay} 秒后启动下一个服务...")
                time.sleep(self.delay)
        
        print("\n" + "=" * 50)
        print(f"批量启动完成！")
        print(f"成功启动: {success_count} 个服务")
        print(f"启动失败: {failed_count} 个服务")
        
        if self.launched_screens:
            print(f"\n已启动的screen会话:")
            for session in self.launched_screens:
                print(f"  - {session}")
            print(f"\n查看所有screen会话: screen -list")
            print(f"连接到会话: screen -r <会话名>")
            print(f"分离会话: Ctrl+A, D")
        
        return failed_count == 0
    
    def list_running_sessions(self):
        """列出所有正在运行的相关screen会话"""
        try:
            result = subprocess.run(['screen', '-list'], 
                                  capture_output=True, text=True)
            lines = result.stdout.split('\n')
            
            service_sessions = []
            for line in lines:
                if 'service_' in line:
                    service_sessions.append(line.strip())
            
            if service_sessions:
                print("正在运行的服务screen会话:")
                for session in service_sessions:
                    print(f"  {session}")
            else:
                print("没有找到正在运行的服务screen会话")
                
        except Exception as e:
            print(f"查询screen会话时发生错误: {e}")
    
    def kill_all_service_sessions(self):
        """终止所有服务相关的screen会话"""
        try:
            result = subprocess.run(['screen', '-list'], 
                                  capture_output=True, text=True)
            lines = result.stdout.split('\n')
            
            killed_count = 0
            for line in lines:
                if 'service_' in line:
                    # 提取会话名称
                    parts = line.strip().split('\t')
                    if len(parts) > 0:
                        session_id = parts[0].split('.')[0]
                        try:
                            subprocess.run(['screen', '-S', session_id, '-X', 'quit'], 
                                         check=True)
                            print(f"已终止screen会话: {session_id}")
                            killed_count += 1
                        except subprocess.CalledProcessError:
                            print(f"终止screen会话失败: {session_id}")
            
            print(f"总共终止了 {killed_count} 个服务screen会话")
            
        except Exception as e:
            print(f"终止screen会话时发生错误: {e}")


def main():
    parser = argparse.ArgumentParser(description='批量启动ROS服务')
    parser.add_argument('-c', '--config', 
                       help='配置文件路径 (默认: ../config/config.yaml)')
    parser.add_argument('-d', '--delay', type=int, default=2,
                       help='服务启动间隔时间（秒，默认: 2）')
    parser.add_argument('-l', '--list', action='store_true',
                       help='列出正在运行的服务screen会话')
    parser.add_argument('-k', '--kill', action='store_true',
                       help='终止所有服务相关的screen会话')
    
    args = parser.parse_args()
    
    launcher = BatchServiceLauncher(config_file=args.config, delay=args.delay)
    
    if args.list:
        launcher.list_running_sessions()
    elif args.kill:
        launcher.kill_all_service_sessions()
    else:
        success = launcher.launch_all_services()
        sys.exit(0 if success else 1)


if __name__ == '__main__':
    main() 